// Usage: ./cvStereoEstimation [-d] -depth d_image -conf c_image -o g_image image_left image_right
// NEED TO DEFINE kernel_width;

#include "hw3.h"
#include <limits.h>
#define kernel_width 9

int display; 
char *leftInputImageFileName;
char *rightInputImageFileName;
char *depthImageFileName;
char *confImageFileName;
char *grayImageFileName;

int focal_length = 10; // mm
int translation = 500; // mm
int disparity;
float depth;

int parseCommandLineArgs(int argc, char** argv)
{

	static struct option long_options[] =
	{
		//{name, has_arg, int *flag, int val}
		{"depth", 1, 0, 0},
		{"conf", 1, 0, 0},
		{0,0,0,0}
	};

	display = 0;
	//depthImageFileName = "\0";
	//confImageFileName = "\0";
	//grayImageFileName = "\0";
	//leftInputImageFileName = "\0";
	//rightInputImageFileName = "\0";

	int option_index = 0;
	int c;
	while (1) {
		c = getopt_long (argc, argv, "do:",
		       long_options, &option_index);

		if (c == -1)
			break;

		switch (c) {
			case 'd':
				display = 1;
			break;
			case 'o':
				grayImageFileName = optarg;
			break;
			break;
			case '?':
	  			/* getopt_long already printed an error message. */
	  		break;
			// long only options
			case 0:
			switch (option_index) {
				case 0:
					depthImageFileName = optarg; 
				break;
				case 1:
					confImageFileName = optarg;
				break;
			}
			break;
		}
	}

	if (optind < argc) {
		// the unspecified option is the input file name
		leftInputImageFileName = argv[optind++];
		rightInputImageFileName = argv[optind];
	}
	return 1;
} 

float comp_std_dev(int *arr, int n)
{
	int       sum = 0;           /* The sum of all elements so far. */
	int       sum_sqrs = 0;      /* The sum of their squares. */

	for (int i = 0; i < n; i++)
	{
		sum += arr[i];
		sum_sqrs += arr[i]*arr[i];
	}
	float average = (float)sum / (float)n;
	float variance = (float)sum_sqrs / (float)n - (average)*(average);

	return sqrt(variance);
}


int main(int argc, char *argv[])
{
	parseCommandLineArgs(argc, argv);
	
	IplImage* lImg = cvLoadImage(leftInputImageFileName);
	IplImage* rImg = cvLoadImage(rightInputImageFileName);
	
	//cvNamedWindow("lImg", CV_WINDOW_AUTOSIZE);
	//cvShowImage("lImg", lImg);
	//cvWaitKey();
	//cvNamedWindow("rImg", CV_WINDOW_AUTOSIZE);
	//cvShowImage("rImg", rImg);
	//cvWaitKey();
	
	if (!lImg || !rImg)
	{
		printf("Error loading input imagesTEST\n");
		return 0;
	}	
	
	IplImage* lGrayImg = cvCreateImage(cvSize(lImg->width,lImg->height),lImg->depth,1);
	IplImage* rGrayImg = cvCreateImage(cvSize(rImg->width,rImg->height),rImg->depth,1);
	IplImage* dispImg = cvCreateImage(cvSize(lImg->width,lImg->height),lImg->depth,1);
	IplImage* depthImg = cvCreateImage(cvSize(lImg->width,lImg->height),lImg->depth,1);
	IplImage* depthImgScaled = cvCreateImage(cvSize(lImg->width,lImg->height),lImg->depth,1);
	IplImage* confImg = cvCreateImage(cvSize(lImg->width,lImg->height),lImg->depth,1);
	IplImage* confImgScaled = cvCreateImage(cvSize(lImg->width,lImg->height),lImg->depth,1);
	IplImage* grayImg = cvCreateImage(cvSize(lImg->width,lImg->height),lImg->depth,1);
	IplImage* grayImgScaled = cvCreateImage(cvSize(lImg->width,lImg->height),lImg->depth,1);
	
	cvCvtColor(lImg,lGrayImg,CV_BGR2GRAY);
	cvCvtColor(rImg,rGrayImg,CV_BGR2GRAY);	
	
	//cvNamedWindow("lGrayImg", CV_WINDOW_AUTOSIZE);
	//cvShowImage("lGrayImg", lGrayImg);
	//cvWaitKey();
	//cvNamedWindow("rGrayImg", CV_WINDOW_AUTOSIZE);
	//cvShowImage("rGrayImg", rGrayImg);
	//cvWaitKey();	
		
	int kernel_offset = (kernel_width - 1) / 2;
	
	for (int row = kernel_offset; row < (lGrayImg->height) - kernel_offset; row++)
	{
		for (int col = kernel_offset; col < (lGrayImg->width) - kernel_offset; col++)
		{
			int kernel[kernel_width][kernel_width];
			for (int krow = 0; krow < kernel_width; krow++)
			{
				for (int kcol = 0; kcol < kernel_width; kcol++)
				{
					kernel[krow][kcol] = ((uchar *)(lGrayImg->imageData + (row - kernel_offset + krow)*lGrayImg->widthStep))[(col - kernel_offset + kcol)];
				}
			}
			
			
			int diff_sqr[kernel_width][kernel_width];						
			int diff_row[(lGrayImg->width) - (kernel_offset * 2)];
			
			for (int i = 0; i < (lGrayImg->width) - (kernel_offset * 2); i++)
			{
				//diff_row[i] = (pow((pow(2.0,(lGrayImg->depth)) - 1),2)) * pow((float)kernel_width,2);
				diff_row[i] = INT_MAX;
			}		
			
			int diff_col;
			for (diff_col = kernel_offset; diff_col < (lGrayImg->width) - kernel_offset; diff_col++)
			{
				/*
				for (int i = 0; i < kernel_width; i++)
				{
					for (int j = 0; j < kernel_width; j++)
					{
						diff_sqr[i][j] = pow(pow(2.0,(lGrayImg->depth)) - 1,2);	
					}
					//printf("\n");
				}
				*/
				
				for (int krow = 0; krow < kernel_width; krow++)
				{
					for (int kcol = 0; kcol < kernel_width; kcol++)
					{						
						
						int diff = (kernel[krow][kcol] - ((uchar *)(rGrayImg->imageData + (row - kernel_offset + krow)*rGrayImg->widthStep))[(diff_col - kernel_offset + kcol)]);

						diff_sqr[krow][kcol]	= diff * diff;
						
					}
				}						
				
				int SSD = 0;
				
				for (int krow = 0; krow < kernel_width; krow++)
				{
					for (int kcol = 0; kcol < kernel_width; kcol++)
					{								
						SSD += diff_sqr[krow][kcol];
					}
				}
				
				diff_row[diff_col - 1] = SSD;																
				
			}


			
			// Set minSSD as large value to compare against
			//int minSSD = (pow((pow(2.0,(lGrayImg->depth)) - 1),2)) * pow((float)kernel_width,2);
			int minSSD = INT_MAX;
			// Initialize minSSD_col to collect the index of the column with the lowest SSD
			int minSSD_col;		
			// For each cell in the scanned row
			
			for (int diff_col = kernel_offset; diff_col < (lGrayImg->width) - kernel_offset; diff_col++)
			{
				// Set minSSD_col to the index of the column with the lowest SSD
				if (diff_row[diff_col - kernel_offset] < minSSD) {
					minSSD_col = (diff_col - kernel_offset) + kernel_offset;
					minSSD = diff_row[diff_col - kernel_offset];
					//printf("minSSD_col = %d\n", minSSD_col);
				}
			}

			
			// Subtract the location of the midpoint of the lImg
			// from the current column under evaluation
			int xL = col - ((lGrayImg->width) / 2);
			
			// Subtract the location of the midpoint of the rImg
			// from the location of the current lowest SSD value			
			int xR = minSSD_col - ((rGrayImg->width) / 2);
			
			// Calculate disparity (per Brown, et al)
			disparity = xL - xR;

			
			// Set disparity value for use in OpenCV IplImage
			CvScalar disparityScalar;
			disparityScalar.val[0] = disparity;
			cvSet2D(dispImg,row,col, disparityScalar);
			
			// Compute depth: Z=T*(f/d) (per Brown, et al)
			depth = (float)translation * ((float)focal_length / (float)disparity);
			// depth = ((float)translation * ((float)focal_length / (float)disparity)) / 1000;
			//printf("depth = %d\n", depth);
			
			// Set depth value for use in OpenCV IplImage
			CvScalar depthScalar;
			depthScalar.val[0] = depth;
			cvSet2D(depthImg,row,col, depthScalar);

			// Compute the confidence image
			int diff_row_sum = 0;
			
			for (int diff_col = kernel_offset; diff_col < (lGrayImg->width) - kernel_offset; diff_col++)
			{
				// sum all of the values of diff_row
				diff_row_sum += diff_row[diff_col - kernel_offset];
				//printf("diff_row_sum = %d\n", diff_row_sum);
			}			
			
			int diff_row_mean = 0;
			diff_row_mean = diff_row_sum / ((lGrayImg->width) - (kernel_offset * 2));
			//printf("diff_row_mean = %d\n", diff_row_mean);
			
			//printf("minSSD = %d\n", minSSD);
			
			int conf_raw = diff_row_mean - minSSD;
			//printf("conf = %d\n", conf);
			
			float std_dev = comp_std_dev(diff_row, ((lGrayImg->width) - (kernel_offset * 2)));
			
			float conf_z = (conf_raw - diff_row_mean) / std_dev;
			
			// Set depth value for use in OpenCV IplImage
			CvScalar confScalar;
			confScalar.val[0] = conf_z;
			cvSet2D(confImg,row,col, confScalar);	
		}
	}		

	cvConvertScale(depthImg, depthImgScaled, 0.001, 0);

	double confMin;
	double confMax;
	cvMinMaxLoc(confImg, &confMin, &confMax);	
	double confShift = 0 - confMin;
	double confScale = 1 / (confMax-confMin);
	cvConvertScale(confImg, confImgScaled, confScale, confShift);
	
	grayImg = cvCloneImage(depthImg);
			
	double grayMin;
	double grayMax;
	cvMinMaxLoc(grayImg, &grayMin, &grayMax);			
	double grayShift = 0 - grayMin;
	double grayScale = 255 / (grayMax-grayMin);
	cvConvertScale(grayImg, grayImgScaled, grayScale, grayShift);

	//cvNamedWindow("Disparity Image", CV_WINDOW_AUTOSIZE);
	//cvShowImage("Disparity Image", dispImg);
	//cvWaitKey();
	//cvNamedWindow("Depth Image", CV_WINDOW_AUTOSIZE);
	//cvShowImage("Depth Image", depthImgScaled);
	//cvWaitKey();	
	if (display) {
		cvNamedWindow("Confidence Image", CV_WINDOW_AUTOSIZE);
		cvShowImage("Confidence Image", confImgScaled);
		cvNamedWindow("Gray Image", CV_WINDOW_AUTOSIZE);
		cvShowImage("Gray Image", grayImgScaled);
		cvWaitKey();	
	}
	
	//cvSaveImage("DisparityImg.png", dispImg, 0);    
	cvSaveImage(depthImageFileName, depthImgScaled, 0);    	
	cvSaveImage(confImageFileName, confImgScaled, 0);    
	cvSaveImage(grayImageFileName, grayImgScaled, 0);    
	/*
	cvSaveImage(dispImageFileName, dispImg, 0);    
	cvSaveImage(depthImageFileName, depthImgScaled, 0); 	
	cvSaveImage(confImageFileName, confImgScaled, 0);    
	cvSaveImage(grayImageFileName, grayImgScaled, 0); 	
	*/
	
	return 0;	
}
